#include "RSA.h"
#include <bitset>
#include <cmath>
#include <iostream>
#include <random>
#include <ctime>

/**
 * 将字符转为二进制表示的字符串
 */
static string binary(char num)
{
    string res;
    char bitMask = 1 << 7;
    for (int i = 0; i < 8; i++) {
        res.append(std::to_string(bitMask & num ? 1 : 0));
        num = num << 1;
    }
    return res;
}

RSA::RSA()
{
}

RSA::~RSA()
{
}

RSA::KeyPair RSA::genKey()
{
    std::default_random_engine eng(time(nullptr));
    // std::default_random_engine eng;
    // std::uniform_int_distribution<uint32_t> rand(0, 0xffffffff); // p, q最多32bits，为n的一半
    std::uniform_int_distribution<uint32_t> rand(0x8000, 0xffff); // p, q最多16Bits 由于蒙哥马利算法的限制，密钥长度无奈又要减半

    // 制造出两个随机的大素数
    do {
        p = rand(eng);
    } while (!isPrime(p));

    do {
        q = rand(eng);
    } while (!isPrime(q));

    n = p * q;
    uint64_t euler = (p - 1) * (q - 1);
    d = inv(e, euler);

    println("e:", e);
    println("p:", p);
    println("q:", q);
    println("n:", n);
    println("l:", euler);
    println("d:", d);

    if ((1.0 * e * d) >= UINT64_MAX) {
        println("e*d=", 1.0 * e * d);
        println("uint64 max=", double(UINT64_MAX));
        println("overflow 64bits");
    } else {
        println("e*d=", e * d);
        println("ed mod euler", (e * d) % euler);
        println("gcd(d,euler)=", gcd(d, euler));
    }

    KeyPair pub_pri;
    std::map<string, string> pub({ { "e", std::to_string(e) }, { "n", std::to_string(n) } });
    std::map<string, string> pri({ { "d", std::to_string(d) }, { "n", std::to_string(n) } });
    pub_pri.insert(std::make_pair("pub", pub));
    pub_pri.insert(std::make_pair("pri", pri));

    return pub_pri;
}

/**
 * 将明文m (假设m是一个小于n的整数)加密为密文C
 * 计算方法为：C = m^e mod n
 */
vector<uint64_t> RSA::encrypt(const string& msg)
{
    // 每4个字符做一次处理，转成一个和密钥长度等长的整数
    vector<uint64_t> res;
    for (int i = 0; i < msg.size(); i += sizeof(uint64_t) / 2) {
        string digitMsg;
        string temp = msg.substr(i, sizeof(uint64_t) / 2);
        for (char& i : temp) {
            digitMsg.append(binary(i));
        }
        println("msg binaryed:", digitMsg);

        std::bitset<sizeof(uint64_t) / 2 * 8> bint(digitMsg);
        println("msg to ull:", bint.to_ullong());

        res.push_back(montgomery(bint.to_ullong(), e, n)); // core
    }

    return res;
}
/**
 * 将密文C解密为明文m，计算方法为：m = C^d mod n
 */
string RSA::decrypt(const vector<uint64_t>& enc)
{
    vector<uint64_t> res;
    string msg;
    for (auto&& each : enc) {
        res.push_back(montgomery(each, d, n));
    }
    for (uint64_t each : res) {
        println("half decrpted:", each);
        string temp;
        do {
            temp.insert(0, 1, char(each & 0xFF));
        } while (each >>= 8);
        msg.append(temp);
    }

    return msg;
}

bool RSA::isPrime(uint64_t num)
{
    for (uint64_t i = 2; i <= sqrt(num); i++) {
        if (num % i == 0)
            return false;
    }

    return true;
}

/**
 * 基础的欧几里得算法求两数最大公约数
 * 关键在于：gcd(a,b) = gcd(b,a mod b)
 */
uint64_t RSA::gcd(uint64_t a, uint64_t b)
{
    return (a % b) ? gcd(b, a % b) : b;
}

/**
 * 拓展的欧几里得算法求两数最大公约数的同时，并求出满足
 * gcd(a,b) = a*x + b*y 中的x, y值。这可用来求乘法逆元。
 */
uint64_t RSA::gcdEx(uint64_t a, uint64_t b, int64_t& x, int64_t& y)
{
    if (a < b)
        return gcdEx(b, a, y, x);
    if (b == 0) {
        x = 1;
        y = 0;
        return a;
    } else {
        int64_t x_;
        uint64_t r = gcdEx(b, a % b, x_, x);
        y = x_ - a / b * x;
        return r;
    }
}

/**
 * 求e关于欧拉函数fyn的乘法逆元d
 */
uint64_t RSA::inv(uint64_t e, uint64_t fyn)
{
    int64_t x, y;
    if (gcdEx(e, fyn, x, y) == 1) {
        // x = x % fyn;  // 注意类型转换，这是个坑，可以不要此句
        return x >= 0 ? x : fyn + x;
    } else {
        return 0;
    }
}

/**
 * 蒙哥马利大整数模幂算法，快速计算 (m ^ e) % n 的值
 */
uint64_t RSA::montgomery(uint64_t m, uint64_t e, uint64_t n)
{
    uint64_t k = 1;
    m %= n;
    while (e > 0) {
        if ((e & 1) != 0)
            k = (k * m) % n;
        m = (m * m) % n; // 注意溢出
        e >>= 1;
    }
    return k;
}